#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jun 27 21:25:27 2018
@project: 天池比赛-A股主板上市公司公告信息抽取
@group: MZH_314
@author: LHQ
"""

import os

from pyltp import Segmentor
from pyltp import Postagger
from pyltp import NamedEntityRecognizer


LTP_DATA_DIR = os.path.join(os.path.dirname(__file__), 'ltp_data_v3.4.0')  # ltp模型目录的路径
cws_model_path = os.path.join(LTP_DATA_DIR, 'cws.model')  # 分词模型路径，模型名称为`cws.model`
pos_model_path = os.path.join(LTP_DATA_DIR, 'pos.model')  # 词性标注模型路径，模型名称为`pos.model`
ner_model_path = os.path.join(LTP_DATA_DIR, 'ner.model')  # 命名实体识别模型路径，模型名称为`pos.model`


class Ner:
    def __init__(self):
        self.segmentor = Segmentor()
        self.postagger = Postagger()
        self.recognizer = NamedEntityRecognizer()
        
        self.segmentor.load(cws_model_path)
        self.postagger.load(pos_model_path)
        self.recognizer.load(ner_model_path)
        
    def prepare_ner(self, sentence):
        words = self.segmentor.segment(sentence)
        postags = self.postagger.postag(words)  # 词性标注
        
        netags = self.recognizer.recognize(words, postags)
        return words, netags
   
    def recog_org(self, sentence):
        words, netags = self.prepare_ner(sentence)
        sbie = {"S":"S-Ni", "B":"B-Ni", "I": "I-Ni", "E": "E-Ni"}
        orgs = self._join(words, netags, sbie)
        return orgs
    
    def _join(self, words, netags, sbie):
        S = sbie["S"]
        B = sbie["B"]
        I = sbie["I"]
        E = sbie["E"]
        
        ents = []
        tmp = []
        flag = False
        for w, t in zip(words, netags):
            if t == S:
                ents.append(w)
                tmp.clear()
                flag = False
                continue
            if not flag and t == B:
                tmp.append(w)
                flag = True
                continue
            if flag and t == I:
                tmp.append(w)
                continue
            if flag and t == E:
                tmp.append(w)
                ent = "".join(tmp)
                ents.append(ent)
                tmp.clear()
                flag = False
                continue
            if flag and t not in (B, I, E):
                tmp.clear()
                flag=False
                continue
        return ents

    def recog_name(self, sentence):
        words, netags = self.prepare_ner(sentence)
        sbie = {"S":"S-Nh", "B":"B-Nh", "I": "I-Nh", "E": "E-Nh"}
        names = self._join(words, netags, sbie)
        return names
        
    def release(self):
        self.segmentor.release()
        self.postagger.release()
        self.recognizer.release()    


  
if __name__ == "__main__":
    ner = Ner()

    s1 = "我在上海林原科技有限公司兼职工作"
    orgs = ner.recog_org(s1)
    print(s1)
    print(orgs)

    print("\n-----------\n")
    s2 = "大股东王二小"
    name = ner.recog_name(s2)
    print(s2)
    print(name)


    ner.release()

